
#include "Rtypes.h"
#include <fstream>
#include <iostream>
#include <string>
#include <stdio.h>
#include <cstdlib>
#include <cmath>
#include <vector>
#include "vector"
#include "TFile.h"
#include "TChain.h"
#include "TTree.h"
#include "TBranch.h"
#include "TLeaf.h"
#include "TCanvas.h"
#include "TH1.h"
#include "TH2.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TMath.h"
#include "TComplex.h"
#include "TLorentzVector.h"
#include "TLegendEntry.h"
#include "TROOT.h"
using namespace std;
#define NEVENT 49
struct particle{
  int pdgId;
  float Px;
  float Py;
  float Pz;
  float m;
  float x;
  float y;
  float z;
  float t;
  int parent_index;
  int child_index;
 // particle parent;
 // particle child;
  bool check;
};
int main()
{
	gROOT->ProcessLine("#include <vector>");
	ifstream parton_init("ana/parton-initial-afterPropagation.dat");
	ifstream parton_hist("ana/parton-collisionsHistory.dat");
	TFile *myfile = TFile::Open("ampt_output.root","RECREATE");
	int mc_n;
	float v2_initial;
	float v2_final;
        vector<float>   mc_pt;
        vector<float>   mc_px;
        vector<float>   mc_py;
        vector<float>   mc_pz;
        vector<float>   mc_m;
        vector<float>   mc_eta;
        vector<float>   mc_phi;
        vector<int>     mc_status;
        vector<int>     mc_pdgId;
        vector<float>   mc_vx_x;
        vector<float>   mc_vx_y;
        vector<float>   mc_vx_z;
        vector<float>   mc_t;
        vector<int>  mc_child_index;
        vector<int>  mc_parent_index;
        vector<int>  N_coll_before;
        vector<int>  N_coll_after;
	float   n_mc_pt;
	float   n_mc_px;
	float   n_mc_py;
	float   n_mc_pz;
	float   n_mc_m;
	float   n_mc_eta;
	float   n_mc_phi;
	int     n_mc_status;
	int     n_mc_pdgId;
	float   n_mc_vx_x;
	float   n_mc_vx_y;
	float   n_mc_vx_z;
	float   n_mc_t;
	float   n_mc_e;
	int     n_N_coll_before;
        int     n_N_coll_after;

	TTree *mytree = new TTree("ampt","parton level from AMPT");
	mytree->Branch("mc_n",&mc_n);
	mytree->Branch("mc_pt",&mc_pt);
	mytree->Branch("mc_px",&mc_px);
	mytree->Branch("mc_py",&mc_py);
	mytree->Branch("mc_pz",&mc_pz);
	mytree->Branch("mc_m",&mc_m);
	mytree->Branch("mc_eta",&mc_eta);
	mytree->Branch("mc_phi",&mc_phi);
	mytree->Branch("mc_vx_x",&mc_vx_x);
	mytree->Branch("mc_vx_y",&mc_vx_y);
	mytree->Branch("mc_vx_z",&mc_vx_z);
	mytree->Branch("mc_t",&mc_t);
	mytree->Branch("mc_status",&mc_status);
	mytree->Branch("mc_pdgId",&mc_pdgId);
	mytree->Branch("mc_child_index",&mc_child_index);
	mytree->Branch("mc_parent_index",&mc_parent_index);
	mytree->Branch("N_coll_before",&N_coll_before);
	mytree->Branch("N_coll_after",&N_coll_after);
	mytree->Branch("v2_initial",&v2_initial);
	mytree->Branch("v2_final",&v2_final);

	int n_init;
	double temp;
	char c = '0';
	int event;
	while(c != '=')
		parton_hist>>c;
	parton_hist >> event >> temp >> temp >> temp;
	particle parton_stg[100000];
	particle array_hist[100000];
	for(int ii = 1;ii<=NEVENT;ii++ )
	{
		cout<<"evnet : "<<ii<<endl;
		long count = 0;
		long count_i = 0;
		long count_f = 0;
		while(event == ii)
		{
			parton_hist>>n_mc_pdgId;
			parton_hist>>n_mc_px>>n_mc_py>>n_mc_pz>>n_mc_m;
			parton_hist>>n_mc_vx_x>>n_mc_vx_y>>n_mc_vx_z>>n_mc_t;
			parton_stg[count].pdgId = n_mc_pdgId;
			parton_stg[count].Px = n_mc_px;
			parton_stg[count].Py = n_mc_py;
			parton_stg[count].Pz = n_mc_pz;
			parton_stg[count].m = n_mc_m;
			parton_stg[count].x = n_mc_vx_x;
			parton_stg[count].y = n_mc_vx_y;
			parton_stg[count].z = n_mc_vx_z;
			parton_stg[count].t = n_mc_t;
			parton_stg[count].parent_index = -1;
			parton_stg[count].child_index = count+2;
			parton_stg[count].check = false;
			count++;

			parton_hist>>n_mc_pdgId;
			parton_hist>>n_mc_px>>n_mc_py>>n_mc_pz>>n_mc_m;
			parton_hist>>n_mc_vx_x>>n_mc_vx_y>>n_mc_vx_z>>n_mc_t;
			parton_stg[count].pdgId = n_mc_pdgId;
			parton_stg[count].Px = n_mc_px;
			parton_stg[count].Py = n_mc_py;
			parton_stg[count].Pz = n_mc_pz;
			parton_stg[count].m = n_mc_m;
			parton_stg[count].x = n_mc_vx_x;
			parton_stg[count].y = n_mc_vx_y;
			parton_stg[count].z = n_mc_vx_z;
			parton_stg[count].t = n_mc_t;
			parton_stg[count].parent_index = -1;
			parton_stg[count].child_index = count+2;
			parton_stg[count].check = false;
			count++;
			
				
			parton_hist>>n_mc_pdgId;
			parton_hist>>n_mc_px>>n_mc_py>>n_mc_pz>>n_mc_m;
			parton_hist>>n_mc_vx_x>>n_mc_vx_y>>n_mc_vx_z>>n_mc_t;
			parton_stg[count].pdgId = n_mc_pdgId;
			parton_stg[count].Px = n_mc_px;
			parton_stg[count].Py = n_mc_py;
			parton_stg[count].Pz = n_mc_pz;
			parton_stg[count].m = n_mc_m;
			parton_stg[count].x = n_mc_vx_x;
			parton_stg[count].y = n_mc_vx_y;
			parton_stg[count].z = n_mc_vx_z;
			parton_stg[count].t = n_mc_t;
			parton_stg[count].parent_index = count-2;
			parton_stg[count].child_index = -1;
			parton_stg[count].check = false;
			count++;


			parton_hist>>n_mc_pdgId;
			parton_hist>>n_mc_px>>n_mc_py>>n_mc_pz>>n_mc_m;
			parton_hist>>n_mc_vx_x>>n_mc_vx_y>>n_mc_vx_z>>n_mc_t;
			parton_stg[count].pdgId = n_mc_pdgId;
			parton_stg[count].Px = n_mc_px;
			parton_stg[count].Py = n_mc_py;
			parton_stg[count].Pz = n_mc_pz;
			parton_stg[count].m = n_mc_m;
			parton_stg[count].x = n_mc_vx_x;
			parton_stg[count].y = n_mc_vx_y;
			parton_stg[count].z = n_mc_vx_z;
			parton_stg[count].t = n_mc_t;
			parton_stg[count].parent_index = count-2;
			parton_stg[count].child_index = -1;
			parton_stg[count].check = false;
			count++;


			c = '0';
			while(c != '=')
				parton_hist>>c;
			parton_hist >> event >> temp >> temp >> temp;
		}
		for(int i=0;i<count;i++)
		{
			if(parton_stg[i].child_index == -1)
			{
				for(int j=i+1;j<count;j++)
				{
					bool same = (fabs(parton_stg[i].Px - parton_stg[j].Px)<0.000001) 
								&& (fabs(parton_stg[i].Py - parton_stg[j].Py)<0.000001)
								&& (fabs(parton_stg[i].Pz - parton_stg[j].Pz)<0.000001)
								&&(fabs(parton_stg[i].x - parton_stg[j].x)<0.000001)
								&&(fabs(parton_stg[i].y - parton_stg[j].y)<0.000001)
								&&(fabs(parton_stg[i].z - parton_stg[j].z)<0.000001)
								&&(parton_stg[i].pdgId == parton_stg[j].pdgId); 
					if(same)
					{
						parton_stg[i].child_index = parton_stg[j].child_index;//the j-th particle will be clear, the index of child will be (index-1)
						parton_stg[parton_stg[j].child_index].parent_index = i;
						for(int k=0;k<count;k++)
						{
							if(k>=j)
								parton_stg[k] = parton_stg[k+1];
							if(parton_stg[k].parent_index > j)
								parton_stg[k].parent_index--;
							if(parton_stg[k].child_index > j)
								parton_stg[k].child_index--;
						}
						count--;
					}
				}
			}
		}
		for(int i=0;i<count;i++)
        	{
            		TLorentzVector v1;
            		n_mc_e = sqrt(pow(parton_stg[i].Px,2)+pow(parton_stg[i].Py,2)+pow(parton_stg[i].Pz,2)+pow(parton_stg[i].m,2));
            		v1.SetPxPyPzE(parton_stg[i].Px,parton_stg[i].Py,parton_stg[i].Pz,n_mc_e);
			
            		mc_pdgId.push_back(parton_stg[i].pdgId);
            		mc_px.push_back(parton_stg[i].Px);
            		mc_py.push_back(parton_stg[i].Py);
            		mc_pz.push_back(parton_stg[i].Pz);
            		mc_vx_x.push_back(parton_stg[i].x);
            		mc_vx_y.push_back(parton_stg[i].y);
            		mc_vx_z.push_back(parton_stg[i].z);
            		mc_m.push_back(parton_stg[i].m);
            		mc_t.push_back(parton_stg[i].t);
            		mc_pt.push_back(v1.Pt());
            		mc_eta.push_back(v1.Eta());
            		mc_phi.push_back(v1.Phi());
			mc_parent_index.push_back(parton_stg[i].parent_index);
			mc_child_index.push_back(parton_stg[i].child_index);
			int mother,daughter;
			n_N_coll_before = 0;
			mother = i;
			while(parton_stg[mother].parent_index>0)
			{
				n_N_coll_before++;
				mother = parton_stg[mother].parent_index;
			}
			n_N_coll_after = 0;
			daughter = i;
			while(parton_stg[daughter].child_index>0)
			{
				n_N_coll_after++;
				daughter = parton_stg[daughter].child_index;
			}

			N_coll_before.push_back(n_N_coll_before);
			N_coll_after.push_back(n_N_coll_after);
			v2_initial = 0;
			v2_final = 0;
			if(n_N_coll_before == 0)
			{	
				count_i++;
				v2_initial =+ cos(v1.Phi());
			}
			if(n_N_coll_after == 0)
			{
				v2_final =+ cos(v1.Phi());
				count_f++;
			}
       	 	}
		int Ncoll;
		float init_px,init_py,init_pz,init_m,init_x,init_y,init_z,init_t,init_e;
		int init_id;
		parton_init>>temp>>temp>>Ncoll>>temp>>temp>>temp>>temp;
		int count_n = 0;
		for(int i=0;i<Ncoll;i++)
		{
			parton_init>>init_id>>init_px>>init_py>>init_pz>>init_m>>init_x>>init_y>>init_z>>init_t;
			bool stored = false;
			for(int j=0;j<count;j++)
			{
				if(fabs(init_px - parton_stg[j].Px)<0.000001
					 && fabs(init_py - parton_stg[j].Py)<0.000001
					 && fabs(init_pz - parton_stg[j].Pz)<0.000001
					)
				{
					stored = true;
					break;
				}
			}
			if(!stored)
			{
				count_n ++;
            			TLorentzVector v1;
            			init_e = sqrt(pow(init_px,2)+pow(init_py,2)+pow(init_pz,2)+pow(init_m,2));
            			v1.SetPxPyPzE(init_px,init_py,init_pz,init_e);
            			mc_pdgId.push_back(init_id);
            			mc_px.push_back(init_px);
            			mc_py.push_back(init_py);
            			mc_pz.push_back(init_pz);
            			mc_vx_x.push_back(init_x);
            			mc_vx_y.push_back(init_y);
            			mc_vx_z.push_back(init_z);
            			mc_t.push_back(init_t);
            			mc_m.push_back(init_m);
            			mc_pt.push_back(v1.Pt());
            			mc_eta.push_back(v1.Eta());
            			mc_phi.push_back(v1.Phi());
				mc_parent_index.push_back(-1);
				mc_child_index.push_back(-1);
				n_N_coll_before = 0;
				n_N_coll_after = 0;
				N_coll_before.push_back(n_N_coll_before);
				N_coll_after.push_back(n_N_coll_after);
				v2_initial =+ cos(v1.Phi());
				v2_final =+ cos(v1.Phi());
				count_i++;
				count_f++;
			}
			
		}
		
		mc_n = Ncoll;
		v2_initial = 1.0*v2_initial/count_i;
		v2_final = 1.0*v2_final/count_f;
		mytree->Fill();
	}
	myfile->Write();
	return 1;
}
